# Please install OpenAI SDK first: `pip3 install openai`
import os
from openai import OpenAI


# 检查API密钥是否存在
# api_key = os.environ.get('DEEPSEEK_API_KEY')
# if not api_key:
#     raise ValueError("请设置DEEPSEEK_API_KEY环境变量")

# client = OpenAI(api_key=api_key, base_url="https://api.deepseek.com")

dashscope_api_key = os.environ.get('DASHSCOPE_API_KEY')
if not dashscope_api_key:
    raise ValueError("请设置DASHSCOPE_API_KEY环境变量")

client = OpenAI(
    api_key=dashscope_api_key,
    base_url="https://dashscope.aliyuncs.com/compatible-mode/v1"
)


# # Round 1
# messages = [{"role": "user", "content": "9.11 and 9.8, which is greater?"}]
# response = client.chat.completions.create(
#     model="deepseek-reasoner",
#     messages=messages,
#     stream=True
# )

# reasoning_content = ""
# content = ""

# for chunk in response:
#     if chunk.choices[0].delta.reasoning_content:
#         reasoning_content += chunk.choices[0].delta.reasoning_content
#     else:
#         content += chunk.choices[0].delta.content or ""

# # 修复打印语句中的格式错误
# print("reasoning_content: {}".format(reasoning_content))
# print("content: {}".format(content))



# ... existing code ...

# 从glob模块导入glob函数，用于文件路径匹配
from glob import glob

# 创建一个空列表，用于存储文本行
text_lines = []


# print(glob("*", recursive=True))

# 遍历所有匹配"milvus_docs/en/faq/*.md"模式的文件路径
# recursive=True允许使用递归模式匹配
for file_path in glob("code/milvus_docs/milvus_docs/en/faq/*.md", recursive=True):
    # 打开当前文件，以只读模式
    with open(file_path, "r") as file:
        # 读取整个文件内容到file_text变量中
        file_text = file.read()

    # 将文件内容按"# "分割，并将结果添加到text_lines列表中
    # split("# ")会根据Markdown中的标题标记分割内容
    text_lines += file_text.split("# ")

# 添加调试信息，检查是否正确读取了文件
print(f"Total text segments: {len(text_lines)}")

# 过滤掉空行和只包含空白字符的行
text_lines = [line.strip() for line in text_lines if line.strip()]
print(f"Non-empty text segments: {len(text_lines)}")

# ... existing code ...

# ... existing code ...

# 修复前的代码存在语法错误：embedding() 应该是 embedding 属性而不是方法调用
# client.embeddings.create(input=text, model="text-embedding-3-small")
# .data[0]
# .embedding()  # 这里有错误
    
# 修复后的正确代码应该是：

def emb_text(text):
    return (
        client.embeddings.create(input=text, model="text-embedding-v2")
        .data[0]
        .embedding  # 访问 embedding 属性而不是调用 embedding() 方法
    )

test_embedding = emb_text("This is a test")
embedding_dim = len(test_embedding)
print(embedding_dim)
print(test_embedding[:10])


from pymilvus import MilvusClient

milvus_client = MilvusClient(uri="./milvus_demo.db")

collection_name = "my_rag_collection"


if milvus_client.has_collection(collection_name):
    milvus_client.drop_collection(collection_name)




# 含义: 向量距离度量类型，这里使用的是"IP"(Inner Product，内积)
# 作用: 定义如何计算向量之间的相似度或距离
# 常见选项:
# "IP": 内积(Inner Product)，适合已归一化的向量，值越大表示越相似
# "L2": 欧几里得距离(Euclidean Distance)，值越小表示越相似
# "COSINE": 余弦相似度，值越大表示越相似(-1到1之间)


milvus_client.create_collection(
    collection_name=collection_name,
    dimension=embedding_dim,
    metric_type="IP",  # Inner product distance
    consistency_level="Bounded",  # Supported values are (`"Strong"`, `"Session"`, `"Bounded"`, `"Eventually"`). See https://milvus.io/docs/consistency.md#Consistency-Level for more details.
)


from tqdm import tqdm

data = []

# 添加调试信息
print(f"Processing {len(text_lines)} text segments...")

for i, line in enumerate(tqdm(text_lines, desc="Creating embeddings")):
    data.append({"id": i, "vector": emb_text(line), "text": line})

print(f"Inserting {len(data)} records into Milvus...")
milvus_client.insert(collection_name=collection_name, data=data)


question = "How is data stored in milvus?"


search_res = milvus_client.search(
    collection_name=collection_name,
    data=[
        emb_text(question)
    ],  # Use the [emb_text](file:///mnt/e/blog/milvus/code/milvus_docs/RAGMilvus.py#L76-L81) function to convert the question to an embedding vector
    limit=3,  # Return top 3 results
    search_params={"metric_type": "IP", "params": {}},  # Inner product distance
    output_fields=["text"],  # Return the text field
)


import json

retrieved_lines_with_distances = [
    (res["entity"]["text"], res["distance"]) for res in search_res[0]
]
print(json.dumps(retrieved_lines_with_distances, indent=4))